#include "stdafx.h"
#include "InstructionMgr.h"
#include "OpCodeMap.h"

InstructionManager::InstructionManager(IByteStreamReader* bytesReader)
: m_bytesReader(bytesReader)
{
    m_instrMap.clear();
}

InstructionManager::~InstructionManager()
{
    m_instrMap.clear();
    if(m_bytesReader != NULL)
    {
        delete m_bytesReader;
        m_bytesReader = NULL;
    }
}

void InstructionManager::GetInstructionList()
{
    if(m_bytesReader == NULL)
        return;
    if(m_bytesReader->GetCurrentBaseAddr() == NULL)
        return;   
    while(true)
    {
        CString instructionStr;    
        DWORD dwInstrAddr = (DWORD)m_bytesReader->GetCurrentBaseAddr();
        if(!GetOneInstruction(instructionStr))
            break;
        m_instrMap[dwInstrAddr] = instructionStr;
        if (instructionStr.CompareNoCase(StrRet) == 0)
            break;
    }
}

bool InstructionManager::GetOneInstruction(CString& instructionStr)
{
    if(m_bytesReader == NULL)
        return false;
    if(m_bytesReader->GetCurrentBaseAddr() == NULL)
        return false;
    // Get Operand
    BYTE opCode = 0;
    m_bytesReader->ReadBytesFromMemAddr(&opCode, sizeof(BYTE));
    // Update EIP address to get following MODRM
    m_bytesReader->UpdateEIPAddress(sizeof(BYTE));
    CString opCodeStr;
    if(!GetOpCode(opCode, opCodeStr))
    {
        CString debugStr;
        debugStr.Format(_T("Unknown Opcode %#x in address %#x\n"), opCode, 
            (DWORD)m_bytesReader->GetCurrentBaseAddr() - sizeof(BYTE));
        OutputDebugString(debugStr);
        return false;
    }
    // if segment override
    CString ptrStrSelection;
    CString segOverrideStr;
    CString ptrOverrideStr = StrEmpty;
    
    if(opCodeStr.CompareNoCase(StrPtrOverride) == 0)
    {
        ptrOverrideStr = StrWORDPTR;
        opCode = 0;
        m_bytesReader->ReadBytesFromMemAddr(&opCode, sizeof(BYTE));
        // Update EIP address to get following opCode
        m_bytesReader->UpdateEIPAddress(sizeof(BYTE));
        if(!GetOpCode(opCode, opCodeStr))
        {
            CString debugStr;
            debugStr.Format(_T("Unknown %s Opcode %#x in address %#x\n"), ptrOverrideStr, 
                opCode, (DWORD)m_bytesReader->GetCurrentBaseAddr() - 2 * sizeof(BYTE));
            OutputDebugString(debugStr);
            return false;
        }
    }

    if(opCodeStr.CompareNoCase(StrFS) == 0)
    {
        segOverrideStr = opCodeStr;
        opCode = 0;
        m_bytesReader->ReadBytesFromMemAddr(&opCode, sizeof(BYTE));
        // Update EIP address to get following opCode
        m_bytesReader->UpdateEIPAddress(sizeof(BYTE));
        if(!GetOpCode(opCode, opCodeStr))
        {
            CString debugStr;
            debugStr.Format(_T("Unknown %s Opcode %#x in address %#x\n"), segOverrideStr, 
                opCode, (DWORD)m_bytesReader->GetCurrentBaseAddr() - 2 * sizeof(BYTE));
            OutputDebugString(debugStr);
            return false;
        }
    }

    // Sub Add need Mod to confirm
    if(opCodeStr.CompareNoCase(StrNeedModConfirm) == 0)
    {
        BYTE ModRM = 0;  
        m_bytesReader->ReadBytesFromMemAddr(&ModRM, sizeof(BYTE));
        // Just read following RegRM byte to confirm operand, no need to update pBaseAddr;
        if(!GetExtendedOpCode(opCode, ModRM, opCodeStr))
        {
            CString debugStr;
            debugStr.Format(_T("Unknown Mod Opcode %#x %#x in address %#x\n"), opCode, ModRM, 
                (DWORD)m_bytesReader->GetCurrentBaseAddr() - sizeof(BYTE));        
            OutputDebugString(debugStr);
            return false;
        }
    }

    // Jcc more operand Code needed
    if(opCodeStr.CompareNoCase(StrNeedMoreOperandCode) == 0)
    {
        BYTE SecOperCode = 0;  
        m_bytesReader->ReadBytesFromMemAddr(&SecOperCode, sizeof(BYTE));
        m_bytesReader->UpdateEIPAddress(sizeof(BYTE));
        if(!GetExtendedOpCode(opCode, SecOperCode, opCodeStr))
        {
            CString debugStr;
            debugStr.Format(_T("Unknown more than one Opcode %#x %#x in address %#x\n"), opCode, SecOperCode, 
            (DWORD)m_bytesReader->GetCurrentBaseAddr() - 2 * sizeof(BYTE));        
            OutputDebugString(debugStr);
            return false;
        }
    }

    if(FindInstructionDNA(opCodeStr, StrRegMem16))
    {
        ptrOverrideStr = StrWORDPTR;
        // Movzx only
        opCodeStr.Replace(StrRegMem16, StrRegMem32);
    }
    if(FindInstructionDNA(opCodeStr, StrReg32) || 
       FindInstructionDNA(opCodeStr, StrReg8) )
    {
        BYTE ModRM = 0;  
        m_bytesReader->ReadBytesFromMemAddr(&ModRM, sizeof(BYTE));
        // if no member effective address needed Update EIP, other wise go on
        if(!FindInstructionDNA(opCodeStr, StrRegMem32) &&
           !FindInstructionDNA(opCodeStr, StrRegMem8) )
            m_bytesReader->UpdateEIPAddress(sizeof(BYTE));
        // Get Reg and EffAddr if needed.
        CString regStr;
        PraseModRMReg(ModRM, regStr);
        if(FindInstructionDNA(opCodeStr,StrReg8))
        {
            // Replace EAX with AL
            if(!MapReg32To8(regStr))
                return false;
            opCodeStr.Replace(StrReg8, regStr); 
        }
        if(!ptrOverrideStr.IsEmpty())
        {
            // Replace EAX with AX
            if(!MapReg32To16(regStr))
                return false;
            opCodeStr.Replace(StrReg32, regStr); 
        }
        opCodeStr.Replace(StrReg32, regStr); 
    }

    if(FindInstructionDNA(opCodeStr, StrRegMem32) ||
       FindInstructionDNA(opCodeStr, StrRegMem8) )
    {   
        BYTE ModRM = 0;  
        m_bytesReader->ReadBytesFromMemAddr(&ModRM, sizeof(BYTE));
        m_bytesReader->UpdateEIPAddress(sizeof(BYTE));

        CString effAddrStr;
        PraseModRMEffAddr(ModRM, effAddrStr);   
        // Get SIB
        if(FindInstructionDNA(effAddrStr, StrSIB) )
        {
            BYTE SIB = 0;  
            m_bytesReader->ReadBytesFromMemAddr(&SIB, sizeof(BYTE));

            m_bytesReader->UpdateEIPAddress(sizeof(BYTE));
            CString strSIBBase;
            PraseSIBBase(SIB, strSIBBase);
            CString strSIBScaleIndex;
            PraseSIBScaleIndex(SIB, strSIBScaleIndex);
            // if not none make base and index together
            if(strSIBScaleIndex.CompareNoCase(StrNone) != 0)
            {
                CString strTemp = StrOpenBracket + strSIBBase + StrPlus;
                strSIBScaleIndex.Replace(StrOpenBracket, strTemp);
                strSIBBase = strSIBScaleIndex;
            }     
            else
            {
                strSIBBase = StrOpenBracket + strSIBBase + StrCloseBracket;
            }
            effAddrStr.Replace(StrSIB, strSIBBase);
        }
        if(FindInstructionDNA(effAddrStr,StrDisp8))
        {
            BYTE Disp8 = 0;  
            m_bytesReader->ReadBytesFromMemAddr(&Disp8, sizeof(BYTE));

            m_bytesReader->UpdateEIPAddress(sizeof(BYTE));
            CString strTemp;
            if ( (Disp8 & 0x80) != 0)
            {
                Disp8 = Disp8 ^ 0xff;
                Disp8 += 1;
                char dispChar[4];
                _itoa_s(Disp8, dispChar, 4, 16);
                strTemp = StrNegetive + CString(dispChar).MakeUpper() + StrCloseBracket;
            }
            else
            {
                char dispChar[4];
                _itoa_s(Disp8, dispChar, 4, 16);
                strTemp = StrPlus + CString(dispChar).MakeUpper() + StrCloseBracket;
            }
            effAddrStr.Replace(StrDisp8, _T(""));
            effAddrStr.Replace(StrCloseBracket, strTemp);
        }
        if(FindInstructionDNA(effAddrStr,StrDisp32))
        {
            DWORD Disp32 = 0;  
            //if(!ptrOverrideStr.IsEmpty())
            //{
            //    WORD Disp16 = 0;
            //    m_bytesReader->ReadBytesFromMemAddr(&Disp16, sizeof(WORD));
            //    m_bytesReader->UpdateEIPAddress(sizeof(WORD));
            //    WORD SigExtHighWord = 0;
            //    WORD SigExtLowWord = 0;
            //    __asm
            //    {
            //        MOV AX, Disp16   
            //        CWD
            //        MOV SigExtHighWord, DX
            //        MOV SigExtLowWord, AX
            //    }
            //    Disp32 = (SigExtHighWord << 16) + SigExtLowWord;;
            //}
            //else
            //{
                m_bytesReader->ReadBytesFromMemAddr(&Disp32, sizeof(DWORD));
                m_bytesReader->UpdateEIPAddress(sizeof(DWORD));
            //}
            // need also deal with negative
            char Disp32Char[12];
            _itoa_s(Disp32, Disp32Char, 12, 16);
            CString tempStr(Disp32Char);
            if(effAddrStr.CompareNoCase(StrDisp32) == 0)
            {
                effAddrStr = StrOpenBracket + effAddrStr + StrCloseBracket;
            }
            else
            {
                effAddrStr.Replace(StrDisp32, StrEmpty);
                if((Disp32 & 0x80000000) != 0)
                {
                    Disp32 = Disp32 ^ 0xffffffff;
                    Disp32 += 1;
                    char NegDisp32Char[12];
                    _itoa_s(Disp32, NegDisp32Char, 12, 16);
                    CString NegTempStr(NegDisp32Char);
                    tempStr = NegTempStr;
                    effAddrStr.Replace(StrCloseBracket, StrNegetive + StrDisp32 + StrCloseBracket);
                }
                else
                   effAddrStr.Replace(StrCloseBracket, StrPlus + StrDisp32 + StrCloseBracket);
            }
            effAddrStr.Replace(StrDisp32, tempStr.MakeUpper());
        }


        // Add dwptr and byte ptr 
        if(FindInstructionDNA(opCodeStr, StrRegMem32))
        {
            ptrStrSelection = StrDWORDPTR;
            opCodeStr.Replace(StrRegMem32, effAddrStr);
        }
        else if(FindInstructionDNA(opCodeStr, StrRegMem8) )
        {
            ptrStrSelection = StrBYTEPTR;
            if(IsReg32(effAddrStr))
            {
                // Replace EAX with AL
                if(!MapReg32To8(effAddrStr))
                    return false;
            }
            opCodeStr.Replace(StrRegMem8, effAddrStr);
        }
        else
            effAddrStr = effAddrStr; // current not supported, may be StrRegMem16;
    }

    if(!ptrOverrideStr.IsEmpty())
    {
        opCodeStr.Replace(StrImm32, StrImm16);
    }

    if(FindInstructionDNA(opCodeStr,StrImm8))
    {
        BYTE Imm8 = 0;  
        m_bytesReader->ReadBytesFromMemAddr(&Imm8, sizeof(BYTE));

        m_bytesReader->UpdateEIPAddress(sizeof(BYTE));
        char Imm8Char[4];
        _itoa_s(Imm8, Imm8Char, 4, 16);
        CString tempStr(Imm8Char);
        opCodeStr.Replace(StrImm8, tempStr.MakeUpper());
    }

    if(FindInstructionDNA(opCodeStr,StrImm16))
    {
        WORD Imm16 = 0;  
        m_bytesReader->ReadBytesFromMemAddr(&Imm16, sizeof(WORD));

        m_bytesReader->UpdateEIPAddress(sizeof(WORD));
        char Imm16Char[10];
        _itoa_s(Imm16, Imm16Char, 10, 16);
        CString tempStr(Imm16Char);
        opCodeStr.Replace(StrImm16, tempStr.MakeUpper());
    }

    if(FindInstructionDNA(opCodeStr,StrImm32))
    {
        DWORD Imm32 = 0;  
        m_bytesReader->ReadBytesFromMemAddr(&Imm32, sizeof(DWORD));

        m_bytesReader->UpdateEIPAddress(sizeof(DWORD));
        char Imm32Char[12];
        _itoa_s(Imm32, Imm32Char, 12, 16);
        CString tempStr(Imm32Char);
        opCodeStr.Replace(StrImm32, tempStr.MakeUpper());
    }

    if(FindInstructionDNA(opCodeStr, StrSegOffset32) )
    {
        DWORD SegOffset32 = 0;  
        m_bytesReader->ReadBytesFromMemAddr(&SegOffset32, sizeof(DWORD));
        m_bytesReader->UpdateEIPAddress(sizeof(DWORD));
        char SetOffset32Char[12];
        _itoa_s(SegOffset32, SetOffset32Char, 12, 16);
        CString tempStr(SetOffset32Char);
        tempStr = StrOpenBracket + tempStr + StrCloseBracket;
        opCodeStr.Replace(StrSegOffset32, tempStr.MakeUpper());    
        ptrStrSelection = StrDWORDPTR;
    }

    // Call near, relative, displacement relative to next instruction
    if(FindInstructionDNA(opCodeStr, StrRel32) )
    {
        DWORD Rel32 = 0;  
        m_bytesReader->ReadBytesFromMemAddr(&Rel32, sizeof(DWORD));

        m_bytesReader->UpdateEIPAddress(sizeof(DWORD));
        // negetive?
        DWORD CallAddress = (DWORD)m_bytesReader->GetCurrentBaseAddr() + Rel32;
        char Rel32Char[12];
        _itoa_s(CallAddress, Rel32Char, 12, 16);
        CString tempStr(Rel32Char);
        opCodeStr.Replace(StrRel32, tempStr.MakeUpper());    
    }

    // JCC near, relative, displacement relative to next instruction
    if(FindInstructionDNA(opCodeStr, StrRel8) )
    {
        BYTE Rel8 = 0;  
        m_bytesReader->ReadBytesFromMemAddr(&Rel8, sizeof(BYTE));
        m_bytesReader->UpdateEIPAddress(sizeof(BYTE));
        DWORD JmpAddress = 0;
        if ( (Rel8 & 0x80) != 0)
        {
            Rel8 = Rel8 ^ 0xff;
            Rel8 += 1;
            JmpAddress = (DWORD)m_bytesReader->GetCurrentBaseAddr() - Rel8;
        }
        else
        {
            JmpAddress = (DWORD)m_bytesReader->GetCurrentBaseAddr() + Rel8;
        }
        char Rel8Char[12];
        _itoa_s(JmpAddress, Rel8Char, 12, 16);
        CString tempStr(Rel8Char);
        opCodeStr.Replace(StrRel8, tempStr.MakeUpper());    
    }
   
    // And SigExt8
    if(FindInstructionDNA(opCodeStr,StrSigExt8))
    {
        BYTE SigExtImm8 = 0;  
        m_bytesReader->ReadBytesFromMemAddr(&SigExtImm8, sizeof(BYTE));
        m_bytesReader->UpdateEIPAddress(sizeof(BYTE));
        DWORD SigExtDWImm8 = 0;
        WORD SigExtHighWord = 0;
        WORD SigExtLowWord = 0;
        __asm
        {
            MOV AL, SigExtImm8
            CBW 
            CWD
            MOV SigExtHighWord, DX
            MOV SigExtLowWord, AX
        }
        SigExtDWImm8 = (SigExtHighWord << 16) + SigExtLowWord;
        char Imm8Char[20];
        _itoa_s(SigExtDWImm8, Imm8Char, 20, 16);
        CString tempStr(Imm8Char);
        opCodeStr.Replace(StrSigExt8, tempStr.MakeUpper());
    }

    if(!ptrOverrideStr.IsEmpty())
    {
        if(!ptrStrSelection.IsEmpty())
        {
            ptrStrSelection = ptrOverrideStr;
        }
    }
    // add segment override and dwptr or byte ptr
    if(!ptrStrSelection.IsEmpty() || !segOverrideStr.IsEmpty() )
    {
        opCodeStr.Replace(StrOpenBracket, ptrStrSelection + segOverrideStr + StrOpenBracket);
    }
    instructionStr = opCodeStr;
    return true;
}